#include "cem/electromagnetics_system.h"
#include "cem/parameter_reader.h"
#include "cem/solver_manager.h"
#include "cem/assembler.h"

#include <libmesh/enum_preconditioner_type.h>
#include <libmesh/enum_solver_type.h>
#include <libmesh/equation_systems.h>
#include <libmesh/linear_implicit_system.h>
#include <libmesh/linear_solver.h>
#include <libmesh/numeric_vector.h>
#include <libmesh/sparse_matrix.h>

#include <chrono>
#include <iostream>

namespace cem
{

// 内部状态：简单记录求解参数
namespace
{
struct SolveParams
{
    double tol{1e-8};
    int max_it{1000};
};
} // namespace

static SolveParams g_params; // 简易全局存储（脚手架阶段）

// 函数: configure_from_parameters
// 作用: 从参数读取容忍度与最大迭代，若未提供则保持默认
void
SolverManager::configure_from_parameters(const ParameterReader & params)
{
    if (auto t = params.get_solver_tolerance())
        g_params.tol = *t;
    if (auto m = params.get_solver_max_iterations())
        g_params.max_it = *m;
}

// 函数: solve
// 作用: 调用 libMesh 的 LinearSolver 接口，使用 GMRES + ILU/AMG 预条件器进行鲁棒求解。
//       注意: 当 libMesh/PETSc 以复数标量构建时，HYPRE BoomerAMG 不支持复数，将导致错误；
//             因此在该情形下强制使用 ILU 作为预条件器以避免崩溃。
void
SolverManager::solve(ElectromagneticsSystem & system)
{
    auto & lin = system.linear_system();

    // 允许系统根据阈值选择求解流程（此处仅打印参数）。
    std::cout << "solver: tol=" << g_params.tol << ", max_it=" << g_params.max_it << std::endl;

    // 在求解之前输出矩阵与右端项的信息，便于排查数值异常
    auto * mat = lin.matrix; // 系统矩阵（可能为 nullptr，取决于装配阶段）
    auto * rhs = lin.rhs;    // 右端项向量

    if (mat)
    {
        std::cout << "solver: matrix info: rows=" << mat->m() << ", cols=" << mat->n()
                  << ", is_closed=" << (mat->closed() ? "true" : "false") << std::endl;
    }
    else
    {
        std::cout << "solver: matrix is null" << std::endl;
    }

    double rhs_l2_norm = 0.0, rhs_linf_norm = 0.0;
    if (rhs)
    {
        rhs_l2_norm = rhs->l2_norm();
        rhs_linf_norm = rhs->linfty_norm();
        std::cout << "solver: rhs info: size=" << rhs->size() << ", l2_norm=" << rhs_l2_norm
                  << ", linfty_norm=" << rhs_linf_norm << std::endl;
    }
    else
    {
        std::cout << "solver: rhs is null" << std::endl;
    }

    // 若右端项为严格零，则直接将解清零并跳过求解，以避免奇异矩阵引发不确定解
    if (rhs && rhs_l2_norm == 0.0 && rhs_linf_norm == 0.0)
    {
        if (lin.solution)
        {
            lin.solution->zero();
            lin.solution->close();
            std::cout << "solver: rhs is zero; set solution to zero and skip solving." << std::endl;

            // 解向量摘要（基于范数）
            const double sol_l2 = lin.solution->l2_norm();
            const double sol_linf = lin.solution->linfty_norm();
            std::cout << "solver: solution stats: size=" << lin.solution->size()
                      << ", l2_norm=" << sol_l2 << ", linfty_norm=" << sol_linf << std::endl;
            return;
        }
    }

    // 使用 GMRES 的鲁棒线性求解流程
    std::cout << "solver: start solving..." << std::endl;
    const auto t_start = std::chrono::high_resolution_clock::now();

    // 构建线性求解器
    auto linear_solver =
        libMesh::LinearSolver<libMesh::Number>::build(lin.get_equation_systems().comm());
    // 选择 GMRES，并根据标量类型/可用性选择合适的预条件器
    linear_solver->set_solver_type(libMesh::GMRES);

#if defined(LIBMESH_USE_COMPLEX_NUMBERS)
    // PETSc/ Hypre 在复数标量下的 AMG（HYPRE BoomerAMG）不受支持，直接改用 ILU
    linear_solver->set_preconditioner_type(libMesh::ASM_PRECOND);
    std::cout
        << "solver: using preconditioner = ASM_PRECOND (complex scalar build; avoiding HYPRE AMG)"
        << std::endl;
#else
    {
        try
        {
            linear_solver->set_preconditioner_type(libMesh::AMG_PRECOND);
            std::cout << "solver: using preconditioner = AMG_PRECOND" << std::endl;
        }
        catch (...)
        {
            // 若 AMG 不可用，则退回到 JACOBI
            linear_solver->set_preconditioner_type(libMesh::JACOBI_PRECOND);
            std::cout << "solver: using preconditioner = JACOBI_PRECOND (AMG unavailable)" << std::endl;
        }
    }
#endif

    linear_solver->init();

    // 确保解向量存在并清零
    if (lin.solution)
    {
        lin.solution->zero();
    }

    // 执行求解（Ax = b）
    auto solve_result =
        linear_solver->solve(*mat, *lin.solution, *rhs, g_params.tol, g_params.max_it);
    unsigned int iters = solve_result.first;
    double final_residual = solve_result.second;

    const auto t_end = std::chrono::high_resolution_clock::now();
    const double elapsed_seconds =
        std::chrono::duration_cast<std::chrono::duration<double>>(t_end - t_start).count();

    // 输出求解摘要
    std::cout << "solver: solve completed, elapsed=" << elapsed_seconds << " s"
              << ", iters=" << iters << ", final_residual=" << final_residual << std::endl;

    // 解向量统计与非数检查（以范数为主，避免对复数 NumericVector 调用min/max）
    if (lin.solution)
    {
        lin.solution->close();
        const double sol_l2 = lin.solution->l2_norm();
        const double sol_linf = lin.solution->linfty_norm();
        std::cout << "solver: solution stats: size=" << lin.solution->size()
                  << ", l2_norm=" << sol_l2 << ", linfty_norm=" << sol_linf << std::endl;
        const bool bad = !(std::isfinite(sol_l2) && std::isfinite(sol_linf));
        if (bad)
        {
            std::cout << "solver WARNING: solution contains non-finite values (NaN/Inf)."
                      << std::endl;
        }
    }
    else
    {
        std::cout << "solver: solution vector is null" << std::endl;
    }
}

// 函数: solve_dual_polarizations_reuse_matrix
// 作用: 组装一次矩阵，针对两套极化（仅边界方向/幅值不同）分别装配 RHS 并使用同一线性求解器两次求解。
//       这样可减少预处理/因子化的重复构建时间，逼近“保留因子分解”的效果。
// 说明: 仅在 rank 0 汇总输出向量，其它 rank 返回空结果。
void
SolverManager::solve_dual_polarizations_reuse_matrix(ElectromagneticsSystem & system,
                                                     const std::vector<BoundarySpec> & specs_pol1,
                                                     const std::vector<BoundarySpec> & specs_pol2,
                                                     std::vector<std::complex<double>> & out_sol_pol1,
                                                     std::vector<std::complex<double>> & out_sol_pol2)
{
    auto & lin = system.linear_system();
    auto & es = system.equation_systems();
    const int rank = es.get_mesh().comm().rank();

    // 1) 首次装配矩阵（基于 system 内部的 boundary_specs）
    {
        Assembler assembler;
        assembler.assemble(system);
    }

    // 2) 构建线性求解器一次，配置 GMRES + 预条件器
    auto linear_solver =
        libMesh::LinearSolver<libMesh::Number>::build(lin.get_equation_systems().comm());
    linear_solver->set_solver_type(libMesh::GMRES);
#if defined(LIBMESH_USE_COMPLEX_NUMBERS)
    linear_solver->set_preconditioner_type(libMesh::ASM_PRECOND);
    if (rank == 0)
        std::cout << "[solver(dual)] using preconditioner = ASM_PRECOND (complex scalar build)"
                  << std::endl;
#else
    try
    {
        linear_solver->set_preconditioner_type(libMesh::AMG_PRECOND);
        if (rank == 0)
            std::cout << "[solver(dual)] using preconditioner = AMG_PRECOND" << std::endl;
    }
    catch (...)
    {
        linear_solver->set_preconditioner_type(libMesh::JACOBI_PRECOND);
        if (rank == 0)
            std::cout << "[solver(dual)] using preconditioner = JACOBI_PRECOND (AMG unavailable)"
                      << std::endl;
    }
#endif
    linear_solver->init();

    auto solve_once = [&](const std::vector<BoundarySpec> & specs,
                          std::vector<std::complex<double>> & out_sol) {
        // 2.1) 仅装配 RHS（不改动矩阵）
        {
            Assembler assembler;
            assembler.assemble_rhs_with_specs(system, specs);
        }

        // 2.2) 若 RHS 近似为零，则直接置解零并跳过求解
        double rhs_l2 = lin.rhs ? lin.rhs->l2_norm() : 0.0;
        double rhs_linf = lin.rhs ? lin.rhs->linfty_norm() : 0.0;
        if (lin.solution)
            lin.solution->zero();

        if (rhs_l2 == 0.0 && rhs_linf == 0.0)
        {
            if (lin.solution)
            {
                lin.solution->zero();
                lin.solution->close();
            }
        }
        else
        {
            auto result = linear_solver->solve(*lin.matrix, *lin.solution, *lin.rhs, g_params.tol, g_params.max_it);
            if (rank == 0)
            {
                std::cout << "[solver(dual)] iters=" << result.first
                          << ", final_residual=" << result.second << std::endl;
            }
            lin.solution->close();
        }

        // 2.3) 聚合到 rank 0 的 std::vector<std::complex<double>>
        // 注意: localize_to_one 是一个集体通信调用，必须在所有进程上调用，
        // 否则会导致并行执行时卡死（deadlock）。
        out_sol.clear();
        if (lin.solution)
        {
            using number_t = libMesh::Number;
            std::vector<number_t> gathered_all;
            // collective on all ranks
            lin.solution->localize_to_one(gathered_all, 0);
            if (rank == 0)
            {
                out_sol.reserve(gathered_all.size());
        #if defined(LIBMESH_USE_COMPLEX_NUMBERS)
                for (const auto & z : gathered_all)
                {
                    out_sol.push_back(z);
                }
        #else
                for (const auto & v : gathered_all)
                {
                    out_sol.emplace_back(v, 0.0);
                }
        #endif
            }
        }
    };

    // 3) 两次求解：极化1、极化2
    solve_once(specs_pol1, out_sol_pol1);
    solve_once(specs_pol2, out_sol_pol2);
}

 } // namespace cem